import model as M
import argparse
import PIL
import numpy as np
from tensorflow.keras.preprocessing.image import load_img
from tensorflow.keras.preprocessing.image import array_to_img
from tensorflow.keras.preprocessing.image import img_to_array


def upscale_image(model, img):
    """Predict the result based on input image and restore the image as RGB."""
    ycbcr = img.convert("YCbCr")
    y, cb, cr = ycbcr.split()
    y = img_to_array(y)
    y = y.astype("float32") / 255.0

    input = np.expand_dims(y, axis=0)
    out = model.predict(input)

    out_img_y = out[0]
    out_img_y *= 255.0

    # Restore the image in RGB color space.
    out_img_y = out_img_y.clip(0, 255)
    out_img_y = out_img_y.reshape((np.shape(out_img_y)[0], np.shape(out_img_y)[1]))
    out_img_y = PIL.Image.fromarray(np.uint8(out_img_y), mode="L")
    out_img_cb = cb.resize(out_img_y.size, PIL.Image.BICUBIC)
    out_img_cr = cr.resize(out_img_y.size, PIL.Image.BICUBIC)
    out_img = PIL.Image.merge("YCbCr", (out_img_y, out_img_cb, out_img_cr)).convert(
        "RGB"
    )
    return out_img


def apply(input, output, scale):
    checkpoint_filepath = "tmp/checkpoint"
    model = M.get_model()
    model.load_weights(checkpoint_filepath)

    org_img = load_img(input)
    scaled_img = upscale_image(model, org_img)
    if scale < 3.:
        factor = scale / 3.
        scaled_img = scaled_img.resize((int(scaled_img.size[0] * factor), int(scaled_img.size[1] * factor)),
                                       PIL.Image.BICUBIC)
    scaled_img.save(output)

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('-input', type=str, required=True, help='input image')
    parser.add_argument('-scale', type=float, default=3., help='scale factor[1,3]')
    parser.add_argument('-output', type=str, default='output.jpg', help='output image, default: output.jpg')
    args = parser.parse_args()
    sf = args.scale
    output_img = args.output
    input_img = args.input

    if sf < 1:
        sf = 1.
    elif sf > 3:
        sf = 3.

    apply(input_img, output_img, sf)